# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import logging
from typing import List, Union

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase


class BinanceFuturesSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    subscription_list = []
    for symbol in self._symbol:
      symbol = symbol.lower()
      subscription_list += [
          '%s@depth@100ms' % symbol,
          '%s@aggTrade' % symbol,
          '%s@markPrice' % symbol,
          '%s@ticker' % symbol,
          '%s@bookTicker' % symbol
      ]
    subscription_str = '/'.join(subscription_list)
    self._register_url('wss://fstream.binance.com/stream?streams=%s' % subscription_str)
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('binance_ws_futures_depth_%s%s' %
                                            (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('binance_ws_futures_trade_%s%s' %
                                            (group_id, worker_id))[0]
    topic_mark_price = topic_map.search_topic_id('binance_ws_futures_mark_price_%s%s' %
                                                 (group_id, worker_id))[0]
    topic_ticker = topic_map.search_topic_id('binance_ws_futures_ticker_%s%s' %
                                             (group_id, worker_id))[0]
    topic_book_ticker = topic_map.search_topic_id('binance_ws_futures_book_ticker_%s%s' %
                                                  (group_id, worker_id))[0]

    self._publisher = {
        'depth@100ms': queue.get_writer(topic_id=topic_depth[1]),
        'aggTrade': queue.get_writer(topic_id=topic_trade[1]),
        'markPrice': queue.get_writer(topic_id=topic_mark_price[1]),
        'ticker': queue.get_writer(topic_id=topic_ticker[1]),
        'bookTicker': queue.get_writer(topic_id=topic_book_ticker[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'binance_ws_topic_event',
                                       [
                                           topic_depth[0],
                                           topic_trade[0],
                                           topic_mark_price[0],
                                           topic_ticker[0],
                                           topic_book_ticker[0]
                                       ])
    self._add_teardown_callback(self._teardown_publisher)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: dict):
    if self.stopped:
      return
    stream = message.get('stream', 'none@none')
    event_type = stream.split('@', maxsplit=1)[1]
    try:
      self._publisher[event_type].write(timestamp, message)
    except KeyError:
      self._logger.error('Unknown event_type: %s' % event_type)
